# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of Cybertron package.
#
# The Cybertron is open-source software based on the AI-framework:
# PyTorch (https://pytorch.org/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Extra activation function
"""

import torch
from torch import nn
from typing import Optional, Union, Type

__all__ = [
    "ShiftedSoftplus",
    "get_activation",
]


class ShiftedSoftplus(nn.Module):
    r"""Compute shifted soft-plus activation function.

    .. math::
       y = \ln\left(1 + e^{-x}\right) - \ln(2)

    Args:
        x (torch.Tensor): input tensor.

    Returns:
        torch.Tensor: shifted soft-plus of input.
    """

    def __init__(self):
        super().__init__()
        self.log2 = torch.log(torch.tensor(2.0))

    def __str__(self):
        return 'ShiftedSoftplus<>'

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.log1p(torch.exp(x)) - self.log2.to(x.device)


_ACTIVATIONS_BY_KEY = {
    'ssp': ShiftedSoftplus,
    'relu': nn.ReLU,
    'relu6': nn.ReLU6,
    'elu': nn.ELU,
    'sigmoid': nn.Sigmoid,
    'tanh': nn.Tanh,
    'leaky_relu': nn.LeakyReLU,
    'prelu': nn.PReLU,
    'gelu': nn.GELU,
    'silu': nn.SiLU,
    'sigmoid': nn.Sigmoid,
    'mish': nn.Mish,
}

_ACTIVATIONS_BY_NAME = {a.__name__: a for a in _ACTIVATIONS_BY_KEY.values()}


def get_activation(cls_name: Optional[Union[str, nn.Module]], **kwargs) -> Optional[nn.Module]:
    """
    Gets the activation function.

    Args:
        cls_name (Optional[Union[str, nn.Module]]): The name or instance of the activation function.
        **kwargs: Additional arguments to pass to the activation function.

    Returns:
        Optional[nn.Module]: The activation function module or None.

    Examples:
        >>> sigmoid = get_activation('sigmoid')
        >>> print(sigmoid)
        Sigmoid()
    """
    if cls_name is None:
        return None

    if isinstance(cls_name, nn.Module):
        return cls_name

    if isinstance(cls_name, str):
        if cls_name.lower() == 'none':
            return None
        if cls_name.lower() in _ACTIVATIONS_BY_KEY:
            return _ACTIVATIONS_BY_KEY[cls_name.lower()](**kwargs)
        if cls_name in _ACTIVATIONS_BY_NAME:
            return _ACTIVATIONS_BY_NAME[cls_name](**kwargs)
        raise ValueError(f"Activation function '{cls_name}' not found")

    raise TypeError(f"Unsupported activation type: {type(cls_name)}")
